package com.tianji.chat.service.impl;

import com.tianji.chat.constants.ApiKeys;
import com.tianji.chat.domain.dto.ChatDTO;
import com.tianji.chat.service.IChatHistoryService;
import com.tianji.common.utils.BeanUtils;
import com.tianji.common.utils.UserContext;
import com.tianji.chat.domain.po.ChatHistory;
import com.tianji.chat.enums.ChatStatus;
import com.tianji.chat.mapper.ChatHistoryMapper;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.RequiredArgsConstructor;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static dev.langchain4j.model.openai.OpenAiModelName.GPT_3_5_TURBO;
import static java.util.stream.Collectors.joining;

/**
 * <p>
 * 聊天记录表 服务实现类
 * </p>
 *
 * @author author
 * @since 2023-12-28
 */
@Service
@RequiredArgsConstructor
public class ChatHistoryServiceImpl extends ServiceImpl<ChatHistoryMapper, ChatHistory> implements IChatHistoryService {

    private final ResourceLoader resourceLoader;

    private String answer_from_chat(String query){
        return "answer from gpt";
    }

    private String answer_from_knowledge(String query) {
        return "from knowledge";
    }


    public String answerMyquestion() {
        // Load the document that includes the information you'd like to "chat" about with the model.
        Resource resource = resourceLoader.getResource("classpath:/file.txt");
//        Resource resource = resourceLoader.getResource("http://example.com/file.txt");
        Path filePath;
        try {
            filePath = Paths.get(resource.getFile().getPath());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        Document document = FileSystemDocumentLoader.loadDocument(filePath);
        System.out.println(document);

        // Split document into segments 100 tokens each
        DocumentSplitter splitter = DocumentSplitters.recursive(
                100,
                0,
                new OpenAiTokenizer(GPT_3_5_TURBO)
        );
        List<TextSegment> segments = splitter.split(document);

        // Embed segments (convert them into vectors that represent the meaning) using embedding model
        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

        // Store embeddings into embedding store for further search / retrieval
        EmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        embeddingStore.addAll(embeddings, segments);

        // Specify the question you want to ask the model
        String question = "Who is Charlie?";

        // Embed the question
        Embedding questionEmbedding = embeddingModel.embed(question).content();

        // Find relevant embeddings in embedding store by semantic similarity
        // You can play with parameters below to find a sweet spot for your specific use case
        int maxResults = 3;
        double minScore = 0.7;
        List<EmbeddingMatch<TextSegment>> relevantEmbeddings
                = embeddingStore.findRelevant(questionEmbedding, maxResults, minScore);

        // Create a prompt for the model that includes question and relevant embeddings
        PromptTemplate promptTemplate = PromptTemplate.from(
                "Answer the following question to the best of your ability:\n"
                        + "\n"
                        + "Question:\n"
                        + "{{question}}\n"
                        + "\n"
                        + "Base your answer on the following information:\n"
                        + "{{information}}");

        String information = relevantEmbeddings.stream()
                .map(match -> match.embedded().text())
                .collect(joining("\n\n"));

        Map<String, Object> variables = new HashMap<>();
        variables.put("question", question);
        variables.put("information", information);

        Prompt prompt = promptTemplate.apply(variables);

        // Send the prompt to the OpenAI chat model
        ChatLanguageModel chatModel = OpenAiChatModel.withApiKey(ApiKeys.OPENAI_API_KEY);
        AiMessage aiMessage = chatModel.generate(prompt.toUserMessage()).content();

        // See an answer from the model
        String answer = aiMessage.text();
        System.out.println(answer); // Charlie is a cheerful carrot living in VeggieVille...
        return answer;
    }

    @Override
    public void saveReply(ChatDTO ChatDTO) {
        // 1.获取登录用户
        Long userId = UserContext.getUser();
        // 2.获取回答
        String answer="";
//        ChatStatus status = ChatDTO.getStatus();
//        if (status==ChatStatus.CHAT){
//            answer=answer_from_chat(ChatDTO.getQuery());
//        }else {
//            answer=answer_from_knowledge(ChatDTO.getQuery());
//        }
        // 2.5 数据封装
        ChatHistory question = BeanUtils.copyBean(ChatDTO, ChatHistory.class);
        question.setUserId(userId);
        question.setResponse(answer);
        // 3.写入数据库
        save(question);

    }


}
